"""
输入：rx, ry, rz, tx, ty, tz
输出：各数据的统计表格，误差最大的DRR图像信息。
备注：drr_visualization必须在draw_histogram之前被调用，否则会出错
"""
import pandas as pd
from pyecharts.charts import Scatter3D
from utils.painter import *
import utils.image_process as imp
from  utils.label_transform import rt2noised_rt
import numpy as np
import cv2 as cv


# 计算图片相似度
def caculate_NCC(image1, image2):
    return np.mean(np.multiply((image1 - np.mean(image1)), (image2 - np.mean(image2)))) \
        / (np.std(image1) * np.std(image2))


class ResAnalyse:
    def __init__(self, data_generator=None, saving_path=None):
        """
        :param data_generator: 用于生成DRR
        :param saving_path: 结果保存的位置
        """
        self.data_generator = data_generator
        # 文件保存路径
        self.saving_path = saving_path
        # 真实值
        self.tru_rx = np.array([0])
        self.tru_ry = np.array([0])
        self.tru_rz = np.array([0])
        self.tru_tx = np.array([0])
        self.tru_ty = np.array([0])
        self.tru_tz = np.array([0])
        # 预测值
        self.pre_rx = np.array([0])
        self.pre_ry = np.array([0])
        self.pre_rz = np.array([0])
        self.pre_tx = np.array([0])
        self.pre_ty = np.array([0])
        self.pre_tz = np.array([0])
        # 最大顶点距离
        self.VM = [0]

    def load_data(self, tru, pre, VM):
        """
        导入非迭代数据，pre的大小和tru的大小一致
        :param tru: 形如 pd.DataFrame(columns=['tru_rx', 'tru_ry', 'tru_rz', 'tru_tx', 'tru_ty', 'tru_tz'])
        :param pre: 形如 pd.DataFrame(columns=['pre_rx', 'pre_ry', 'pre_rz', 'pre_tx', 'pre_ty', 'pre_tz'])
        :param VM: 最大顶点距离
        :return:
        """
        # 真实值
        self.tru_rx = np.array(tru['tru_rx'])
        self.tru_ry = np.array(tru['tru_ry'])
        self.tru_rz = np.array(tru['tru_rz'])
        self.tru_tx = np.array(tru['tru_tx'])
        self.tru_ty = np.array(tru['tru_ty'])
        self.tru_tz = np.array(tru['tru_tz'])
        # 预测值
        self.pre_rx = np.array(pre['pre_rx'])
        self.pre_ry = np.array(pre['pre_ry'])
        self.pre_rz = np.array(pre['pre_rz'])
        self.pre_tx = np.array(pre['pre_tx'])
        self.pre_ty = np.array(pre['pre_ty'])
        self.pre_tz = np.array(pre['pre_tz'])
        # 最大顶点距离
        self.VM = np.array(VM)

    def draw_histogram(self):
        draw_distribution_histogram(data=self.VM, path=self.saving_path + '/VM',
                                    x_label="VM/mm", y_label='frequency', title='VM_error_distribution',
                                    color1='pink', color2='red')
        # 绘制每个维度的差值
        draw_distribution_histogram(data=abs(self.tru_rx - self.pre_rx), path=self.saving_path + '/rx',
                                    x_label="angle", y_label='frequency', title='rx_error_distribution',
                                    color1='yellow', color2='orange')
        draw_distribution_histogram(data=abs(self.tru_tx - self.pre_tx), path=self.saving_path + '/tx',
                                    x_label="t/mm", y_label='frequency', title='tx_error_distribution',
                                    color1='steelblue', color2='blue')
        draw_distribution_histogram(data=abs(self.tru_ry - self.pre_ry), path=self.saving_path + '/ry',
                                    x_label="angle", y_label='frequency', title='ry_error_distribution',
                                    color1='springgreen', color2='green')
        draw_distribution_histogram(data=abs(self.tru_ty - self.pre_ty), path=self.saving_path + '/ty',
                                    x_label="t/mm", y_label='frequency', title='ty_error_distribution',
                                    color1='pink', color2='deeppink')
        draw_distribution_histogram(data=abs(self.tru_rz - self.pre_rz), path=self.saving_path + '/rz',
                                    x_label="angle", y_label='frequency', title='rz_error_distribution',
                                    color1='violet', color2='purple')
        draw_distribution_histogram(data=abs(self.tru_tz - self.pre_tz), path=self.saving_path + '/tz',
                                    x_label="t/mm", y_label='frequency', title='tz_error_distribution',
                                    color1='lightsalmon', color2='salmon')

    def drr_visualization(self, VM_threshold, saving_drr_path=None, posture=None, mode='save_img'):
        """
        给定VM的选取范围，输出真实值和预测值对应的DRR图
        :param VM_threshold: 大于VM阈值的图像会被保存
        :param saving_drr_path: drr的保存路径
        :param posture: "标准正位"/"标准侧位"
        :param mode:是否展示图片，是否保存图片
        :return: 
        """
        for i in range(len(self.VM)):
            if self.VM[i] > VM_threshold:
                r_noise = rt2noised_rt(np.array([self.tru_rx[i], self.tru_ry[i], self.tru_rz[i]]), mode=posture)
                rx_noise, ry_noise, rz_noise = r_noise[0], r_noise[1], r_noise[2]
                tru_drr = imp.get_new_drr_numpy(rx_noise,
                                                ry_noise,
                                                rz_noise,
                                                self.tru_tx[i],
                                                self.tru_ty[i],
                                                self.tru_tz[i],
                                                posture,
                                                self.data_generator)
                pre_drr = imp.get_new_drr_numpy(rx_noise,
                                                ry_noise,
                                                rz_noise,
                                                self.pre_tx[i],
                                                self.pre_ty[i],
                                                self.pre_tz[i],
                                                posture,
                                                self.data_generator)
                tru_r = np.array([self.tru_rx[i], self.tru_ry[i], self.tru_rz[i]])
                tru_t = np.array([self.tru_tx[i], self.tru_ty[i], self.tru_tz[i]])
                pre_r = np.array([self.pre_rx[i], self.pre_ry[i], self.pre_rz[i]])
                pre_t = np.array([self.pre_tx[i], self.pre_ty[i], self.pre_tz[i]])
                plt.figure(figsize=(8, 5))
                plt.subplot(1, 2, 1)
                plt.title('实际值'
                          + "\n" +
                          '角度' + str(np.around(tru_r, 2))
                          + "\n" +
                          '平移' + str(np.around(tru_t, 2)))
                plt.imshow(tru_drr.reshape(128, 128), cmap='gray')
                plt.subplot(1, 2, 2)
                plt.title('预测值'
                          + "\n" +
                          '角度' + str(np.around(pre_r, 2))
                          + "\n" +
                          '平移' + str(np.around(pre_t, 2)))
                plt.imshow(pre_drr.reshape(128, 128), cmap='gray')
                ncc = float(caculate_NCC(tru_drr, pre_drr))
                mse = float(np.mean((tru_drr - pre_drr) ** 2))
                plt.suptitle("第{0}组_NCC:{1}_MSE:{2}".format(i, round(ncc, 3), round(mse, 3)))
                if mode == 'display':
                    plt.show()
                elif mode == 'save_img':
                    plt.savefig(saving_drr_path + "第{}组：VM{}.png".format(i, round(self.VM[i], 5)))
                plt.close()

    def plot_3D_point(self, VM_threshold, saving_drr_path=None):
        """
        给定VM的选取范围,查看预测偏差较大的点和较小的点在采样空间中的分布
        :param VM_threshold: 大于VM阈值的点被视作坏点
        :param saving_drr_path: 保存路径
        :return: 
        """
        good_point_rx = []
        good_point_ry = []
        good_point_rz = []
        good_point_tx = []
        good_point_ty = []
        good_point_tz = []
        bad_point_rx = []
        bad_point_ry = []
        bad_point_rz = []
        bad_point_tx = []
        bad_point_ty = []
        bad_point_tz = []
        for i in range(len(self.VM)):
            if self.VM[i] > VM_threshold:
                bad_point_rx.append(self.tru_rx[i])
                bad_point_ry.append(self.tru_ry[i] - 270)
                bad_point_rz.append(self.tru_rz[i] - 90)
                bad_point_tx.append(self.tru_tx[i])
                bad_point_ty.append(self.tru_ty[i])
                bad_point_tz.append(self.tru_tz[i])
            else:
                good_point_rx.append(self.tru_rx[i])
                good_point_ry.append(self.tru_ry[i] - 270)
                good_point_rz.append(self.tru_rz[i] - 90)
                good_point_tx.append(self.tru_tx[i])
                good_point_ty.append(self.tru_ty[i])
                good_point_tz.append(self.tru_tz[i])
        # print(list(zip(good_point_rx, good_point_ry, good_point_rz)))
        # print(list(zip(bad_point_rx, bad_point_ry, bad_point_rz)))
        # 绘制角度
        draw_3D_scatter(data=[list(zip(good_point_rx, good_point_ry, good_point_rz)),
                              list(zip(bad_point_rx, bad_point_ry, bad_point_rz))],
                        series_name=['good_point_r', 'bad_point_r'],
                        color=['green', 'red'],
                        axis_range=5,
                        page_title='标准正位角度散点分布',
                        title='标准正位(0,270,90)角度散点分布',
                        saving_path=(saving_drr_path + "3D散点角度图.html"))
        # 绘制位移
        draw_3D_scatter(data=[list(zip(good_point_tx, good_point_ty, good_point_tz)),
                              list(zip(bad_point_tx, bad_point_ty, bad_point_tz))],
                        series_name=['good_point_t', 'bad_point_t'],
                        color=['green', 'red'],
                        axis_range=50,
                        page_title='标准正位位移散点分布',
                        title='标准正位±50位移位移散点分布',
                        saving_path=(saving_drr_path + "3D散点位移图.html"))

    def to_csv(self, datas, names, saving_path):
        for data, name in zip(datas, names):
            data.to_csv(saving_path + name)

